package com.ruoyi.ai.ws.handler;

import org.springframework.stereotype.Component;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketSession;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.core.scheduler.Schedulers;

import java.util.concurrent.ConcurrentHashMap;

@Component
public class ChatWebSocketHandler implements WebSocketHandler {

    private final static ConcurrentHashMap<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    public static void sendMessage(String username, String vue) {
        WebSocketSession webSocketSession = sessions.get(username);
        if (webSocketSession != null) {
            webSocketSession.send(Mono.just(webSocketSession.textMessage(vue))).subscribe();
        }
    }

    @Override
    public Mono<Void> handle(WebSocketSession session) {
        // 获取握手信息中的 URI
        String uri = session.getHandshakeInfo().getUri().toString();
        // 提取路径中的 userid
        String userId = extractUserIdFromUri(uri);
        System.out.println("WebSocket connection opened for user: " + userId);
        sessions.put(userId, session);
        // 接收客户端发送的消息
        return session.receive()
                .publishOn(Schedulers.boundedElastic())
                .doOnNext(msg -> {
                    String payload = msg.getPayloadAsText();
                    System.out.println("Received message: " + payload);
                    // 向客户端发送回显消息
                    session.send(Mono.just(session.textMessage("收到: " + payload)))
                            .subscribe();
                }).doFinally(signalType -> {
                    // 当连接关闭时，移除会话
                    sessions.remove(userId);
                    System.out.println("WebSocket connection closed for user: " + userId);
                })
                .then(); // 等待连接关闭
    }

    /**
     * 从 URI 中提取 userid
     */
    private String extractUserIdFromUri(String uri) {
        // 假设路径格式为 /ws/{userid}
        String[] parts = uri.split("/");
        if (parts.length > 2) {
            return parts[parts.length - 1]; // 返回最后一个部分作为 userid
        }
        return "unknown"; // 默认值
    }
}
